#include "Functions/Draw.h"
#include "Containers/RangePair.h"
#include "Functions/Error.h"
#include "Functions/File.h"
#include "Functions/TerminalIO.h"
#include "RooFit.h"
#include "TArrow.h"
#include "TLatex.h"
#include "TLine.h"
#include "TPaveLabel.h"
#include "TStyle.h"
using namespace std;
using namespace BOSS_Afterburner;

/// Auxiliary function that is used by the more specific `DrawAndSave` functions for `TH1D` and `TH2D`.
void BOSS_Afterburner::Draw::DrawAndSave(TH1& hist, const char* saveas, Option_t* opt,
                                         TString logScale)
{
  hist.Draw(opt);
  SaveCanvas(saveas, gPad, logScale);
}

/// Draw and save a 1D distribution (output folder is determined from `FrameworkSettings.h`).
void BOSS_Afterburner::Draw::DrawAndSave(TH1D& hist, const char* saveas, TString logScale)
{
  DrawAndSave(hist, saveas, "ep", logScale);
}

/// Draw and save a 2D distribution (output folder is determined from `FrameworkSettings.h`).
void BOSS_Afterburner::Draw::DrawAndSave(TH2D& hist, const char* saveas, TString logScale)
{
  DrawAndSave(hist, saveas, "colz", logScale);
}

/// Function that allows you to create and save a quick sketch of a `TTree` branch.
void BOSS_Afterburner::Draw::DrawAndSave(TTree* tree, const char* varexp, const char* selection,
                                         Option_t* opt, TString logScale)
{
  tree->Draw(varexp, selection, opt);
  SaveCanvas(Form("%s_%s", tree->GetName(), varexp), gPad, logScale);
}

/// Helper function that allows you to save a `TPad` or `TCanvas`.
void BOSS_Afterburner::Draw::SaveCanvas(const char* saveas, TVirtualPad* pad, TString logScale)
{
  if(TerminalIO::IsEmptyPtr(pad)) return;
  SetLogScale(logScale, pad);
  pad->SaveAs(File::SetOutputFilename(saveas));
}

/// *TEMPORARY* function that serves as a fix for the bug that causes the wrong best pair to be stored.
/// @todo Consider removing `BOSS_Afterburner::Draw::DrawDifference`.
void BOSS_Afterburner::Draw::DrawDifference(TH1* histToDraw, TH1* histToSubtract, Option_t* opt,
                                            const char* setLog)
{
  if(!gPad) return;
  gPad->Clear();
  histToDraw->Scale(-1.);
  histToDraw->Add(histToSubtract);
  histToSubtract->SetLineColor(kWhite);
  // histToSubtract->SetMarkerColor(kWhite);
  TString option(opt);
  if(dynamic_cast<TH1F*>(histToDraw))
  {
    histToSubtract->Draw(option.Data());
    option += "";
  }
  histToDraw->Draw(option.Data());
  BOSS_Afterburner::Draw::SaveCanvas(Form("%s", histToDraw->GetName()), gPad, setLog);
}

/// Helper function that allows you to set the log scale of a `TPad` or `TCanvas`.
void BOSS_Afterburner::Draw::SetLogScale(TString logScale, TVirtualPad* pad)
{
  // * Check arguments * //
  if(BOSS_Afterburner::TerminalIO::IsEmptyPtr(pad)) return;
  logScale.ToLower();
  // * Set log scale * //
  pad->SetLogx(logScale.Contains("x"));
  pad->SetLogy(logScale.Contains("y"));
  pad->SetLogz(logScale.Contains("z"));
}

void BOSS_Afterburner::Draw::IndicateFitRange(const YAML::Node& conf)
{
  const auto& fitdef = conf["Ranges"]["fit"];
  const auto& draw   = conf["Settings"]["Indicate fit range"];
  if(!fitdef) return;
  if(!draw) return;
  if(!draw.as<Bool_t>()) return;

  RangePair fit = fitdef;

  auto line = new TLine;
  line->SetLineWidth(5);
  line->SetLineColor(kGreen);
  SetLineCoordinates(*line, fit.X1(), 0., fit.X2(), 0.);
  line->Draw();
}

void BOSS_Afterburner::Draw::Indicate3SigmaRange(const YAML::Node& conf)
{
  if(!gPad) return;

  const auto& range = conf["3-sigma range"];
  if(!range) return;

  const auto& massDef   = range["mass"];
  const auto& widthDef  = range["width"];
  const auto& sigmasDef = range["sigmas"];
  const auto& fromDef   = range["from"];
  const auto& toDef     = range["to"];
  if(!massDef && !widthDef && !fromDef && !toDef) EXCEPTION("Incomplete key \"3-sigma range\"");

  if(massDef.IsScalar() && (!!widthDef || !!sigmasDef))
  {
    Double_t mass  = massDef.as<Double_t>();
    Double_t width = 0.;
    if(!!widthDef && widthDef.IsScalar())
      width = widthDef.as<Double_t>();
    else if(!!sigmasDef && sigmasDef.IsSequence())
    {
      for(const auto& sigDef : sigmasDef)
        if(sigDef.IsScalar())
        {
          Double_t sigma = sigDef.as<Double_t>();
          sigma *= sigma;
          width += sigma;
        }
      width = sqrt(width);
    }
    Indicate3SigmaRange(mass, width);
  }
  else if(fromDef.IsScalar() && toDef.IsScalar() && massDef.IsScalar())
  {
    Double_t x1 = fromDef.as<Double_t>();
    Double_t x2 = toDef.as<Double_t>();
    Double_t m  = massDef.as<Double_t>();
    Indicate3SigmaRange(x1, m, x2);
  }
}

void BOSS_Afterburner::Draw::Indicate3SigmaRange(Double_t mass, Double_t width)
{
  Indicate3SigmaRange(mass - 3 * width, mass, mass + 3 * width);

  Double_t x1, y1, x2, y2;
  gPad->Update();
  gPad->GetRangeAxis(x1, y1, x2, y2);
  Double_t scaleX = x2 - x1;
  Double_t scaleY = y2 - y1;

  auto a = new TArrow(mass - 3 * width, y2, mass + 3 * width, y2, 0.03, "<>");
  auto t = new TLatex(mass - 3 * width - 0.02 * scaleX, y2, Form("#sigma_{_{tot}} = %g ", width));
  t->SetTextAlign(32);
  t->SetTextAngle(90);
  t->SetTextColor(kGray);
  t->SetTextSize(.03);
  a->SetLineColor(kGray);
  t->Draw();
  a->Draw();
}

void BOSS_Afterburner::Draw::Indicate3SigmaRange(Double_t from, Double_t mid, Double_t to)
{
  if(!gPad) return;

  auto left = new TLine;
  left->SetLineStyle(kDotted);
  left->SetLineColor(kGray + 2);
  auto right  = new TLine(*left);
  auto middle = new TLine(*left);
  middle->SetLineColor(kBlack);
  middle->SetLineWidth(2);

  Double_t x1, y1, x2, y2, m;
  gPad->Update();
  gPad->GetRangeAxis(x1, y1, x2, y2);

  SetLineCoordinates(*left, from, y1, from, y2);
  SetLineCoordinates(*middle, mid, y1, mid, y2);
  SetLineCoordinates(*right, to, y1, to, y2);
  left->Draw();
  middle->Draw();
  right->Draw();
}

void BOSS_Afterburner::Draw::IndicateMassWindow(const YAML::Node& conf)
{
  if(!gPad) return;

  const auto& range = conf["Mass window"];
  if(!range || !range.IsSequence() || range.size() != 2) return;

  auto left = new TLine;
  left->SetLineStyle(kDotted);
  left->SetLineColor(kGray + 2);
  auto right = new TLine(*left);

  Double_t x1, y1, x2, y2;
  gPad->Update();
  gPad->GetRangeAxis(x1, y1, x2, y2);
  x1 = range[0].as<Double_t>();
  x2 = range[1].as<Double_t>();

  SetLineCoordinates(*left, x1, y1, x1, y2);
  SetLineCoordinates(*right, x2, y1, x2, y2);

  left->Draw();
  right->Draw();
}

void BOSS_Afterburner::Draw::SetLineCoordinates(TLine& line, Double_t x1, Double_t y1, Double_t x2,
                                                Double_t y2)
{
  line.SetX1(x1);
  line.SetX2(x2);
  line.SetY1(y1);
  line.SetY2(y2);
}

void BOSS_Afterburner::Draw::PlotStatsOnFrame(RooAbsData& obj, RooPlot* frame)
{
  if(!gStyle->GetOptStat()) return;
  if(!frame) return;

  using namespace RooFit;
  obj.statOn(frame, Layout(0.25, 0.55, .32), Color(kBlue + 2), Format("NU", AutoPrecision(0)));
  auto att = frame->getAttText();
  if(!att) return;
  att->SetTextSize(.035);
  att->SetTextColor(kBlue + 2);
}

void BOSS_Afterburner::Draw::FormatLegend(TLegend* legend, const YAML::Node& conf)
{
  if(!legend) return;
  const auto& def = conf["Legend definition"];
  if(!def.IsMap()) return;

  const auto& posDef = def["position"];
  if(posDef.IsSequence() && posDef.size() == 4)
  {
    legend->SetX1(posDef[0].as<Double_t>());
    legend->SetY1(posDef[1].as<Double_t>());
    legend->SetX2(posDef[2].as<Double_t>());
    legend->SetY2(posDef[3].as<Double_t>());
  }
  if(posDef.IsMap())
  {
    const auto& x1Def = posDef["x1"];
    const auto& x2Def = posDef["x2"];
    const auto& y1Def = posDef["y1"];
    const auto& y2Def = posDef["y2"];
    if(x1Def.IsScalar() && x2Def.IsScalar() && y1Def.IsScalar() && y2Def.IsScalar())
    {
      legend->SetX1(x1Def.as<Double_t>());
      legend->SetY1(y1Def.as<Double_t>());
      legend->SetX2(x2Def.as<Double_t>());
      legend->SetY2(y2Def.as<Double_t>());
    }
  }

  const auto& sizeDef = def["font size"];
  if(sizeDef.IsScalar()) legend->SetTextSize(sizeDef.as<Double_t>());
  if(gPad) gPad->RecursiveRemove(legend);
  legend->Draw();
}

void BOSS_Afterburner::Draw::SetCanvasTitle(const char* title)
{
  if(!title) return;
  gStyle->SetOptTitle(0);
  auto cTitle = new TPaveLabel(.11, .95, .35, .99, title, "brndc");
  cTitle->Draw();
  if(gPad) gPad->Update();
}